import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
import yaml
import os
from dataloader.dataloader import get_dataloader

def save_prediction_to_nc(pred_tensor, config, filename_prefix="prediction_"):
    """
    保存预测或观测结果为 NetCDF 文件，文件名自动编号防止覆盖。
    pred_tensor: torch.Tensor [B, T, C, H, W]
    """
    pred = pred_tensor.cpu().numpy()
    print(f"pred:",pred.shape)
    B, T, C, H, W = pred.shape

    var_names = config.get("var_names", [f"var{i}" for i in range(C)])
    if len(var_names) != C:
        raise ValueError(f"[ERROR] var_names 数量({len(var_names)})与模型输出通道数({C})不一致")

    time = np.arange(T)
    height = np.arange(H)
    width = np.arange(W)

    save_dir = config.get("save_dir", "results/")
    os.makedirs(save_dir, exist_ok=True)

    # 自动编号保存文件
    existing_files = [f for f in os.listdir(save_dir) if f.startswith(filename_prefix) and f.endswith(".nc")]
    idx = len(existing_files)
    save_path = os.path.join(save_dir, f"{filename_prefix}{idx:03d}.nc")

    data_vars = {
        var_names[c]: (("time", "height", "width"), pred[0, :, c, :, :])
        for c in range(C)
    }

    ds = xr.Dataset(
        data_vars=data_vars,
        coords={
            "time": time,
            "height": height,
            "width": width
        }
    )
    ds.to_netcdf(save_path)
    print(f"[INFO] Saved to {save_path}")
    return save_path

def export_test_ground_truth_to_nc(config):
    """
    将测试集的 ground truth y 全部导出为 .nc 文件（一个 batch 样本一份）
    """
    test_loader, test_set = get_dataloader(config, mode="test")
    save_dir = config.get("obs_save_dir", "ground_truth/")
    var_names = config.get("var_names", ["var0"])
    os.makedirs(save_dir, exist_ok=True)

    sample_idx = 0  # 全局样本计数，防止文件名重复

    for _, y in test_loader:
        if y.dim() == 4:
            y = y.unsqueeze(1)  # 保证是5维: [B, T, C, H, W]

        y = test_set.denormalize(y)  # 反归一，y 仍是 tensor

        # 遍历batch维度
        for b in range(y.shape[0]):
            single_sample = y[b:b+1]  # 保持batch维度为1，[1, T, C, H, W]
            prefix = f"obs_{sample_idx:03d}_"
            save_prediction_to_nc(single_sample, {"save_dir": save_dir, "var_names": var_names}, filename_prefix=prefix)
            sample_idx += 1

def compare_nc_variables(pred_path, obs_path, var_names, time_idx=0, pangu_path=None):
    pred = xr.open_dataset(pred_path)
    obs = xr.open_dataset(obs_path)
    pangu = xr.open_dataset(pangu_path) if pangu_path else None

    for var_name in var_names:
        if var_name not in pred or var_name not in obs:
            print(f"[WARN] {var_name} not found in prediction or observation dataset.")
            continue

        pred_var = pred[var_name].isel(time=time_idx)
        obs_var = obs[var_name].isel(time=time_idx)
        pangu_var = None
        if pangu is not None and var_name in pangu:
            pangu_var = pangu[var_name].isel(time=time_idx)

        fig, axs = plt.subplots(1, 3 if pangu_var is not None else 2, figsize=(15, 4))

        pred_var.plot(ax=axs[0], cmap="viridis")
        axs[0].set_title(f"Your Prediction: {var_name}")

        if pangu_var is not None:
            pangu_var.plot(ax=axs[1], cmap="viridis")
            axs[1].set_title(f"PanGu Prediction: {var_name}")
            obs_ax = axs[2]
        else:
            obs_ax = axs[1]

        obs_var.plot(ax=obs_ax, cmap="viridis")
        obs_ax.set_title(f"Observation: {var_name}")

        plt.suptitle(f"{var_name} @ Time {time_idx}")
        plt.tight_layout()
        plt.show()


def plot_nc_error_map(pred_path, obs_path, var_names, time_idx=0, pangu_path=None):
    pred = xr.open_dataset(pred_path)
    obs = xr.open_dataset(obs_path)
    pangu = xr.open_dataset(pangu_path) if pangu_path else None

    for var_name in var_names:
        if var_name not in pred or var_name not in obs:
            print(f"[WARN] {var_name} not found in prediction or observation dataset.")
            continue

        error_pred = (pred[var_name].isel(time=time_idx) - obs[var_name].isel(time=time_idx)) ** 2
        error_pred.plot(cmap="hot", figsize=(5, 4))
        plt.title(f"{var_name} Error Map (Pred vs Obs) @ Time {time_idx}")
        plt.tight_layout()
        plt.show()

        if pangu is not None and var_name in pangu:
            error_pangu = (pangu[var_name].isel(time=time_idx) - obs[var_name].isel(time=time_idx)) ** 2
            error_pangu.plot(cmap="hot", figsize=(5, 4))
            plt.title(f"{var_name} Error Map (PanGu vs Obs) @ Time {time_idx}")
            plt.tight_layout()
            plt.show()


def run_visualization(pred_path, obs_path, var_names, time_idx=0, mode="compare", pangu_path=None):
    if mode == "compare":
        compare_nc_variables(pred_path, obs_path, var_names, time_idx, pangu_path)
    elif mode == "error":
        plot_nc_error_map(pred_path, obs_path, var_names, time_idx, pangu_path)
    else:
        raise ValueError(f"[ERROR] 未知模式: {mode}")


if __name__ == "__main__":
    with open("config.yaml", "r", encoding="utf-8") as f:
        cfg = yaml.safe_load(f)

    # vis_cfg = config.get("vis", {})
    # pred_path = vis_cfg.get("pred_path", "results/prediction_000.nc")
    # obs_path = vis_cfg.get("obs_path", "ground_truth/obs_000.nc")
    # pangu_path = vis_cfg.get("pgpred_path", None)  # 🔹 新增字段
    # var_names = config.get("var_names", ["temperature"])
    # time_idx = vis_cfg.get("time_idx", 0)
    # mode = vis_cfg.get("mode", "compare")
    #
    # if mode == "compare":
    #     compare_nc_variables(pred_path, obs_path, var_names, time_idx, pangu_path)
    # elif mode == "error":
    #     plot_nc_error_map(pred_path, obs_path, var_names, time_idx, pangu_path)
    # elif mode == "export":
    #     export_test_ground_truth_to_nc(config)
    # else:
    #     raise ValueError(f"[ERROR] 未知模式: {mode}")
    vis_cfg = cfg.get("vis", {})
    pred_path = vis_cfg.get("pred_path", "results/prediction_000.nc")
    obs_path = vis_cfg.get("obs_path", "ground_truth/obs_000.nc")
    pangu_path = vis_cfg.get("pgpred_path", None)  # 🔹 新增字段
    var_names = cfg.get("var_names", ["temperature"])
    time_idx = vis_cfg.get("time_idx", 0)
    mode = vis_cfg.get("mode", "compare")

    run_visualization(pred_path, obs_path, var_names, time_idx, mode, pangu_path)
